import torch
import torchvision
import torch.nn.functional as F

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt


from PIL import Image, ImageFont, ImageDraw
# 导入中文字体，指定字号
import matplotlib

if __name__ == "__main__":
	# 有 GPU 就用 GPU，没有就用 CPU
	device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
	matplotlib.rc("font", family='SimHei')  # 中文字体
	font = ImageFont.truetype('SimHei.ttf', 32)
	idx_to_labels = np.load('idx_to_labels.npy', allow_pickle=True).item()
	model = torch.load('classfier.pth')
	model = model.eval().to(device)
	from torchvision import transforms
	# 测试集图像预处理-RCTN：缩放、裁剪、转 Tensor、归一化
	test_transform = transforms.Compose([transforms.Resize(256),
									 transforms.CenterCrop(224),
									 transforms.ToTensor(),
									 transforms.Normalize(
										 mean=[0.485, 0.456, 0.406],
										 std=[0.229, 0.224, 0.225])
									])
	from PIL import Image
	img_path = 'D:\\1.jpg'
	# img_path = 'test_img/banana1.jpg'
	# img_path = 'test_img/test_bananan.jpg'
	# img_path = 'test_img/test_orange.jpg'
	# img_path = 'test_img/test_orange_2.jpg'
	# img_path = 'test_img/test_kiwi.jpg'
	# img_path = 'test_img/test_lemon.jpg'
	# img_path = 'test_img/test_石榴.jpg'
	# img_path = 'test_img/test_火龙果.jpg'
	img_pil = Image.open(img_path)
	input_img = test_transform(img_pil) # 预处理
	input_img = input_img.unsqueeze(0).to(device)
	# 执行前向预测，得到所有类别的 logit 预测分数
	pred_logits = model(input_img)
	pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
	plt.figure(figsize=(22, 10))

	x = idx_to_labels.values()
	y = pred_softmax.cpu().detach().numpy()[0] * 100
	width = 0.45 # 柱状图宽度

	ax = plt.bar(x, y, width)

	plt.bar_label(ax, fmt='%.2f', fontsize=15) # 置信度数值
	plt.tick_params(labelsize=20) # 设置坐标文字大小

	plt.title(img_path, fontsize=30)
	plt.xticks(rotation=45) # 横轴文字旋转
	plt.xlabel('类别', fontsize=20)
	plt.ylabel('置信度', fontsize=20)
	plt.show()
	n = 2
	top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
	pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
	confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度
	draw = ImageDraw.Draw(img_pil)
	for i in range(n):
		class_name = idx_to_labels[pred_ids[i]] # 获取类别名称
		confidence = confs[i] * 100 # 获取置信度
		text = '{:<15} {:>.4f}'.format(class_name, confidence)
		print(text)

	# 文字坐标，中文字符串，字体，rgba颜
		draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1))